#!/usr/bin/env python
import cv2
import numpy as np
from operator import mod
from math import ceil
import sys,os

try:
    mediafile = sys.argv[1]
except:
    mediafile = '1.mp4'
annofile = ''.join([mediafile,'.anno'])

def timing(position):
    current = position//25
    current_sec = mod(current,60)
    current_min = current // 60
    if current_min >= 60:
        current_min = mod(current_min, 60)
    current_hour = current // 3600
    current_time = ''.join([str(current_hour).zfill(2),':',str(current_min).zfill(2),':',str(current_sec).zfill(2)])

    return current_time

def get_last_anno():
    with open(annofile,'r') as f:
        tags = f.read().splitlines()
    index = -1
    if tags[index].isalpha():
        index -= 1 
    end_anno = tags[index]
    if '_' in end_anno:
        end_anno = 'NULL'

    return end_anno

def load_tag_state():
    with open(annofile,'r') as f:
        tags = f.read().splitlines()
    tag_states = [line for line in tags if line.isalpha()]
    try:
        tag_state_current = tag_states[-1]
    except:
        tag_state_current = 'None'

    return tag_state_current

def get_position(time):
    try:
        current_hour,current_min,current_sec = time.split(':')
        current_frame = int(fps * (int(current_hour) * 3600 + int(current_min) * 60 + int(current_sec)))
    except:
        current_frame = position

    return current_frame

def get_number(num_img):
    mean,stddv = cv2.meanStdDev(num_img)
    *_,channels = num_img.shape
    imgMean = 121
    #imgMean = sum(mean)/channels
    _,num_img = cv2.threshold(num_img,imgMean,255,cv2.THRESH_BINARY)
    #_, white_num =cv2.threshold(num_img,121,255,cv2.THRESH_BINARY)
    #inverse_img = 255 - num_img
    #_, black_num =cv2.threshold(inverse_img,121,255,cv2.THRESH_BINARY)
    #new_num_img = cv2.add(white_num,black_num)
    results = [cv2.matchTemplate(num_img,template,cv2.TM_CCOEFF_NORMED).astype(float) for template in nums]

    return mod(results.index(max(results)),10)

def nothing(emp):
    pass

def show_anno(annofile):
    with open(annofile,'r') as f:
        tags = f.read().splitlines()
    print(tags[-1])
    print(tags[-2])

media = cv2.VideoCapture(mediafile)
fps = int(media.get(cv2.CAP_PROP_FPS))
total_s = media.get(cv2.CAP_PROP_FRAME_COUNT)
frames = int(total_s)
total_time = timing(frames)
loop_flag = 0
position = 0
clip_flag = 0
edit_flag = 0

try:
    _,image = media.read()
except:
    sys.exit()

print('fps = ', str(fps))
print("total_s = ", total_s)

title = ''.join([mediafile,'\tSPACE(start)/ENTER(end) for tag, p(ause), f(orward), b(ack), q(uit), z(oom).'])
window_width = 640
window_height = 480

white_zero = cv2.imread('template/white_0.jpg')
white_one = cv2.imread('template/white_1.jpg')
white_two = cv2.imread('template/white_2.jpg')
white_three = cv2.imread('template/white_3.jpg')
white_four = cv2.imread('template/white_4.jpg')
white_five = cv2.imread('template/white_5.jpg')
white_six = cv2.imread('template/white_6.jpg')
white_seven = cv2.imread('template/white_7.jpg')
white_eight = cv2.imread('template/white_8.jpg')
white_night = cv2.imread('template/white_9.jpg')
black_zero = cv2.imread('template/black_0.jpg')
black_one = cv2.imread('template/black_1.jpg')
black_two = cv2.imread('template/black_2.jpg')
black_three = cv2.imread('template/black_3.jpg')
black_four = cv2.imread('template/black_4.jpg')
black_five = cv2.imread('template/black_5.jpg')
black_six = cv2.imread('template/black_6.jpg')
black_seven = cv2.imread('template/black_7.jpg')
black_eight = cv2.imread('template/black_8.jpg')
black_night = cv2.imread('template/black_9.jpg')
nums = [white_zero,white_one,white_two,white_three,white_four,white_five,white_six,white_seven,white_eight,white_night,black_zero,black_one,black_two,black_three,black_four,black_five,black_six,black_seven,black_eight,black_night]

_,img = cv2.threshold(image,121,255,cv2.THRESH_BINARY)
shape = image.shape
height = shape[0]
width = shape[1]

time_font = cv2.FONT_HERSHEY_COMPLEX
tag_font = cv2.FONT_HERSHEY_SIMPLEX
help_font = cv2.FONT_ITALIC
help_info_1 = 'help:  "a" for add; "r" for retrieve.; "=" for next'
help_info_2 = '      <Space> to tag start; <Enter> to tag end;"." to restore '
time_color = (209,80,0,255)
tag_state_color = (209,80,0,255)
help_color_1 = (109,80,50,155)
help_color_2 = (109,80,50,155)

if height == 720:
    font_size = 1
    time_position = (50,30)
    time_video_start_position = (50,640)
    time_bold = 2
    tag_state_position = (400,30)
    tag_state_bold = 2
    tag_position = (600,30)
    tag_color = (209,80,0,255)
    tag_bold = 2
    help_position_1 = (50,680)
    help_bold_1 = 2
    help_position_2 = (50,710)
    help_bold_2 = 2
    serial_position = (50,600)
elif height==1080:
    font_size = 2
    proportion = 1080/720
    time_position = (ceil(50*proportion),ceil(30*proportion))
    time_video_start_position = (ceil(50*proportion),ceil(640*proportion))
    time_bold = 2
    tag_state_position = (ceil(200+400*proportion),ceil(30*proportion))
    tag_state_bold = 2
    tag_position = (ceil(200+600*proportion),ceil(30*proportion))
    tag_color = (209,80,0,255)
    tag_bold = 2
    help_position_1 = (ceil(50*proportion),ceil(680*proportion))
    help_bold_1 = 2
    help_position_2 = (ceil(50*proportion),ceil(710*proportion))
    help_bold_2 = 2
    serial_position = (ceil(50*proportion),ceil(600*proportion))

if height != 720:
    img = cv2.resize(img,(1280,720),interpolation=cv2.INTER_AREA)
try:
    year_1 = str(get_number(img[53:85, 47:70]))
    year_2 = str(get_number(img[53:85, 71:94]))
    year_3 = str(get_number(img[53:85, 96:119]))
    year_4 = str(get_number(img[53:85, 119:142]))
    month_1 = str(get_number(img[53:85, 167:190]))
    month_2 = str(get_number(img[53:85, 191:214]))
    day_1 = str(get_number(img[53:85, 240:263]))
    day_2 = str(get_number(img[53:85, 263:286]))
    hour_1 = str(get_number(img[53:85, 479:502]))
    hour_2 = str(get_number(img[53:85, 503:526]))
    min_1 = str(get_number(img[53:85, 551:574]))
    min_2 = str(get_number(img[53:85, 575:598]))
    sec_1 = str(get_number(img[53:85, 623:646]))
    sec_2 = str(get_number(img[53:85, 647:670]))
    time_video_start = ''.join([year_1,year_2,year_3,year_4,month_1,month_2,day_1,day_2,'_',hour_1,hour_2,':',min_1,min_2,':',sec_1,sec_2])
except:
    time_video_start = '00000000_00:00:00'


def timing(position):
    current = position//25
    current_sec = mod(current,60)
    current_min = current // 60
    if current_min >= 60:
        current_min = mod(current_min, 60)
    current_hour = current // 3600
    current_time = ''.join([str(current_hour).zfill(2),':',str(current_min).zfill(2),':',str(current_sec).zfill(2)])

    return current_time


cv2.namedWindow(title, 0)
cv2.resizeWindow(title, window_width, window_height)
cv2.createTrackbar('frame',title,0,frames,nothing)

frame_index = 0
if media.isOpened():
    while True:
        if loop_flag == position:
            loop_flag = loop_flag + 1
            cv2.setTrackbarPos('frame',title,loop_flag)
        else:
            position = cv2.getTrackbarPos('frame',title)
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)

        current_time = timing(position)
        time_info = ''.join([current_time,'/',total_time])

        tag_state = load_tag_state()
        last_line = get_last_anno()
        try:
            last_start,last_end = last_line.split('|')
        except:
            last_start = last_line.split('|')[0]
        if last_start != 'NULL':
            last_start_frame = get_position(last_start)
        else:
            last_start_frame = position
        if last_start_frame != position:
            clip_flag = 1
            tag_state = load_tag_state()

        wait_key = cv2.waitKeyEx(1)
        if wait_key == 112: # p
            cv2.waitKey(0)
        elif wait_key == 103: # g
            position = loop_flag + 250
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 102: # f
            position = loop_flag + 125
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 100: # d
            position = loop_flag + 25
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 98: # b
            position = loop_flag - 100
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 118: # v
            position = loop_flag - 25
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 110: # n
            position = loop_flag - 250
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 108: # l
            position = frames - 100
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 122: # z
            window_width = ceil(window_width * 1.2)
            window_height = ceil(window_height * 1.2)
            window_zoom = (window_width,window_height)
            cv2.resizeWindow(title, window_zoom)
        elif wait_key == 120: # x
            window_width = ceil(window_width * 0.8)
            window_height = ceil(window_height * 0.8)
            window_zoom = (window_width,window_height)
            cv2.resizeWindow(title, window_zoom)
        elif wait_key == 8: # backspace
            cmd = ''.join(["sed -i '$d' ",annofile])
            os.system(cmd)
        elif wait_key == 27: # ESC
            cv2.destroyAllWindows()
            sys.exit()
            break
        elif wait_key == 97: # a
            edit_flag = 1
            tag_state = 'add'
            with open(annofile,'a') as f:
                f.writelines('add\n');
        elif wait_key == 114: # r
            edit_flag = 1
            tag_state = 'retrieve'
            with open(annofile,'a') as f:
                f.writelines('retrieve\n');
        elif wait_key == 117: # u
            edit_flag = 1
            tag_state = 'xx'
            with open(annofile,'a') as f:
                f.writelines('xx\n');
            position = frames
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 107: # k
            edit_flag = 1
            tag_state = 'ok'
            with open(annofile,'a') as f:
                f.writelines('ok\n');
            position = frames
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 32: # Space
            clip_flag = 1
            edit_flag = 1
            current_frame = cv2.getTrackbarPos('frame',title)
            start_time = timing(current_frame - 10)
            with open(annofile,'a') as f:
                f.writelines(''.join([start_time,'|']));
        elif wait_key == 13: # Enter
            clip_flag = 1
            edit_flag = 1
            current_frame = cv2.getTrackbarPos('frame',title)
            end_time = timing(current_frame)
            with open(annofile,'a') as f:
                f.writelines(''.join([end_time,'\n']));
        elif wait_key == 46: # .
            clip_flag = 1
            tag_state = load_tag_state()
            last_line = get_last_anno()
            try:
                last_start,last_end = last_line.split('|')
            except:
                last_start = last_line.split('|')[0]
            if last_start != 'NULL':
                last_start_frame = get_position(last_start)
            else:
                last_start_frame = position
            position = last_start_frame
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 47: # /
            clip_flag = 1
            tag_state = load_tag_state()
            last_line = get_last_anno()
            try:
                last_start,last_end = last_line.split('|')
            except:
                last_end = last_line.split('|')[0]
            if last_end != 'NULL':
                last_end_frame = get_position(last_end)
            else:
                last_end_frame = position
            position = last_end_frame
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 45: # -
            position = 1
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        elif wait_key == 61: # =
            #if not load_tag_state() == 'None':
            #    position = frames
            #    loop_flag = position
            #    media.set(cv2.CAP_PROP_POS_FRAMES,position)
            position = frames
            loop_flag = position
            media.set(cv2.CAP_PROP_POS_FRAMES,position)
        success, frame = media.read()
        if success:
            frame_index += 1
            cv2.putText(frame,time_info,time_position,time_font,font_size,time_color,time_bold)
            cv2.putText(frame,help_info_1,help_position_1,help_font,font_size,help_color_1,help_bold_1)
            cv2.putText(frame,help_info_2,help_position_2,help_font,font_size,help_color_2,help_bold_2)
            cv2.putText(frame,''.join(['video start: ',time_video_start]),time_video_start_position,time_font,font_size,help_color_1,help_bold_1)
            if 'tag_state' in locals().keys():
                cv2.putText(frame,tag_state,tag_state_position,tag_font,font_size,tag_state_color,tag_state_bold)
            cv2.imshow(title, frame)
            if clip_flag == 1:
                last_tag = get_last_anno()
                cv2.putText(frame,last_tag,tag_position,tag_font,font_size,tag_color,tag_bold)
        else:
            cv2.destroyWindow(title)
            if edit_flag == 1:
                cmd = ''.join(["vim ",annofile])
                #os.system(cmd)
            if 'tag_state' in locals().keys():
                if tag_state == 'xx':
                    mediafile_useless = Path(mediafile)
                    mediafile_useless.replace(''.join([mediafile,'.bak']))
                if tag_state == 'ok':
                    mediafile_useless = Path(mediafile)
                    mediafile_useless.replace(''.join([mediafile,'.ok']))
            break
    print("frame_index = ", frame_index)
